#! /usr/bin/env python
# coding=utf-8


# 过滤 ALICUT 比对的基因 去除 低低于某个长度的片段

import os

from Bio import SeqIO
import shutil


import sys
import argparse


parser = argparse.ArgumentParser(
    description=''' 对alicut的结果进行过滤 提取出
    用法:
    filter_alicut_alignment_genes.py -i Orthogroup_Sequences_mafft_raw -o Orthogroup_Sequences_mafft_raw_filtered -l 35
    由大天才于2021年7月20日创建于浙江农业大学''')



parser.add_argument('-i',
                help='必须给定，ALICUT的结果文件路径')

parser.add_argument('-o',
                help='必须给定，输出文件路径')

parser.add_argument('-l',
                help='过滤长度 默认为35 对核酸 建议取到 105')

parser.add_argument('-t',
                help='用于找到ALICUT的文件，默认为 ALICUT_')


args = parser.parse_args()

if not args.i or not args.o:
    parser.print_help()
    sys.exit()

infile = args.i


outfile = args.o



if not args.l:

	len_cut_off = 35

else:

	len_cut_off = int(args.l)


if not args.t:

	title = 'ALICUT_'

else:

	title = args.t


try:
	shutil.rmtree(outfile)
except:
	pass

try:
	os.mkdir(outfile)
except:
	pass





for i in os.listdir(infile):

	t = 0

	if i.split('.')[-1] in ['fasta','fa','faa','fna'] and i.find(title)!=-1:
	
		for j in SeqIO.parse(infile+'/'+i,'fasta'):
		
			t = 1
		
			if len(j) < len_cut_off:
				t = 0
				print('filtered %s' %i)
			break

	if t==1:
		os.popen("cp %s %s" % (infile+'/'+i, outfile+'/'+i))
